#include "napi/native_api.h"
#include "hilog/log.h"
#include <dlib/dnn.h>
#include <dlib/gui_widgets.h>
#include <dlib/clustering.h>
#include <dlib/string.h>
#include <dlib/image_io.h>
#include <dlib/image_processing/frontal_face_detector.h>
#include <multimedia/image_framework/image_pixel_map_mdk.h>
#include "facedetectcnn.h"

//define the buffer size. Do not change the size!
//0x9000 = 1024 * (16 * 2 + 4), detect 1024 face at most
#define DETECT_BUFFER_SIZE 0x9000

#undef LOG_DOMAIN
#undef LOG_TAG
#define LOG_DOMAIN 0x0000  // 全局domain宏，标识业务领域
#define LOG_TAG "face"   // 全局tag宏，标识模块日志tag

using namespace std;
using namespace dlib;

template <template <int,template<typename>class,int,typename> class block, int N, template<typename>class BN, typename SUBNET>
using residual = add_prev1<block<N,BN,1,tag1<SUBNET>>>;

template <template <int,template<typename>class,int,typename> class block, int N, template<typename>class BN, typename SUBNET>
using residual_down = add_prev2<avg_pool<2,2,2,2,skip1<tag2<block<N,BN,2,tag1<SUBNET>>>>>>;

template <int N, template <typename> class BN, int stride, typename SUBNET>
using block  = BN<con<N,3,3,1,1,relu<BN<con<N,3,3,stride,stride,SUBNET>>>>>;

template <int N, typename SUBNET> using ares      = relu<residual<block,N,affine,SUBNET>>;
template <int N, typename SUBNET> using ares_down = relu<residual_down<block,N,affine,SUBNET>>;

template <typename SUBNET> using alevel0 = ares_down<256,SUBNET>;
template <typename SUBNET> using alevel1 = ares<256,ares<256,ares_down<256,SUBNET>>>;
template <typename SUBNET> using alevel2 = ares<128,ares<128,ares_down<128,SUBNET>>>;
template <typename SUBNET> using alevel3 = ares<64,ares<64,ares<64,ares_down<64,SUBNET>>>>;
template <typename SUBNET> using alevel4 = ares<32,ares<32,ares<32,SUBNET>>>;

using anet_type = loss_metric<fc_no_bias<128,avg_pool_everything<
                            alevel0<
                            alevel1<
                            alevel2<
                            alevel3<
                            alevel4<
                            max_pool<3,3,2,2,relu<affine<con<32,7,7,2,2,
                            input_rgb_image_sized<150>
                            >>>>>>>>>>>>;

// frontal_face_detector detector;
shape_predictor sp;
anet_type net;

void rgba_to_bgr(const unsigned char* rgba, unsigned char* bgr, int width, int height) {
    const int rgba_step = width * 4;
    const int bgr_step = width * 3;
    
    for (int y = 0; y < height; ++y) {
        const unsigned char* rgba_row = rgba + y * rgba_step;
        unsigned char* bgr_row = bgr + y * bgr_step;
        
        for (int x = 0; x < width; ++x) {
            // RGBA: [R,G,B,A] → BGR: [B,G,R]
            bgr_row[0] = rgba_row[2];  // B
            bgr_row[1] = rgba_row[1];  // G
            bgr_row[2] = rgba_row[0];  // R
            
            rgba_row += 4;
            bgr_row += 3;
        }
    }
}

void rgba_to_dlib_grayscale(const unsigned char* rgba_data, int width, int height, dlib::array2d<unsigned char>& dlib_image)
{
    dlib_image.set_size(height, width);
    
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            const int offset = (y * width + x) * 4;
            // 使用简单的亮度公式转换为灰度
            dlib_image[y][x] = static_cast<unsigned char>(
                0.299 * rgba_data[offset] +      // R
                0.587 * rgba_data[offset + 1] +  // G
                0.114 * rgba_data[offset + 2]    // B
            );
        }
    }
}

void rgba_to_dlib_rgb(const unsigned char* rgba_data, int width, int height, dlib::array2d<dlib::rgb_pixel>& dlib_image)
{
    dlib_image.set_size(height, width);
    
    for (int y = 0; y < height; ++y) {
        for (int x = 0; x < width; ++x) {
            const int offset = (y * width + x) * 4;
            dlib_image[y][x].red = rgba_data[offset];
            dlib_image[y][x].green = rgba_data[offset + 1];
            dlib_image[y][x].blue = rgba_data[offset + 2];
            // 忽略 alpha 通道 rgba_data[offset + 3]
        }
    }
}

// 非线性映射函数（指数衰减）
static double distanceToScoreNonlinear(double distance, double scale = 100) {
    double score = scale / (1 + distance);
    if (score < 0) return 0;
    if (score > scale) return scale;
    return score;
}

static napi_value NAPI_Global_feature(napi_env env, napi_callback_info info) {
    // 获取并解析传参
    size_t argc = 2;
    napi_value args[2] = {nullptr};
    napi_get_cb_info(env, info, &argc, args , nullptr, nullptr);

    int _left = 0;
    int _top = 0;
    int _right = 0;
    int _bottom = 0;
    
    napi_value left = nullptr;
    napi_create_string_utf8(env, "left", NAPI_AUTO_LENGTH, &left);
    napi_value left_v = nullptr;
    napi_get_property(env, args[1], left, &left_v);
    napi_get_value_int32(env, left_v, &_left);
    
    napi_value top = nullptr;
    napi_create_string_utf8(env, "top", NAPI_AUTO_LENGTH, &top);
    napi_value top_v = nullptr;
    napi_get_property(env, args[1], top, &top_v);
    napi_get_value_int32(env, top_v, &_top);
    
    napi_value right = nullptr;
    napi_create_string_utf8(env, "right", NAPI_AUTO_LENGTH, &right);
    napi_value right_v = nullptr;
    napi_get_property(env, args[1], right, &right_v);
    napi_get_value_int32(env, right_v, &_right);
    
    napi_value bottom = nullptr;
    napi_create_string_utf8(env, "bottom", NAPI_AUTO_LENGTH, &bottom);
    napi_value bottom_v = nullptr;
    napi_get_property(env, args[1], bottom, &bottom_v);
    napi_get_value_int32(env, bottom_v, &_bottom);
    
//     OH_LOG_INFO(LOG_APP, "left=%{public}d,top=%{public}d,right=%{public}d,bottom=%{public}d", _left, _top, _right, _bottom);
    
    dlib::rectangle rect(_left, _top, _right, _bottom);
    
    // 初始化NativePixelMap对象。
    NativePixelMap *native = OH_PixelMap_InitNativePixelMap(env, args[0]);
    if (native == nullptr) {
        return nullptr;
    }

    // 获取图片信息。
    struct OhosPixelMapInfos pixelMapInfo;
    OH_PixelMap_GetImageInfo(native, &pixelMapInfo);

    // 获取PixelMap对象每行字节数。
    int32_t rowBytes;
    OH_PixelMap_GetBytesNumberPerRow(native, &rowBytes);

    // 获取PixelMap对象是否可编辑的状态。
    int32_t editable = 0;
    OH_PixelMap_GetIsEditable(native, &editable);

    // 获取PixelMap对象是否支持Alpha通道。
    int32_t supportAlpha = 0;
    OH_PixelMap_IsSupportAlpha(native, &supportAlpha);
    
    // 获取PixelMap对象数据的内存地址，并锁定该内存。
    void *pixelAddr = nullptr;
    OH_PixelMap_AccessPixels(native, &pixelAddr);
    
    array2d<unsigned char> img;
    rgba_to_dlib_grayscale((unsigned char *)pixelAddr, pixelMapInfo.width, pixelMapInfo.height, img);
    
    // 释放PixelMap对象数据的内存锁。
    OH_PixelMap_UnAccessPixels(native);
    
//     matrix<rgb_pixel> img;
//     load_image(img, "/data/storage/el2/base/haps/entry/files/3.bmp");
    //获取指定一个区域的人脸形状
    dlib::full_object_detection shape = sp(img, rect);

    //计算该人脸的特征值
    std::vector<dlib::matrix<dlib::rgb_pixel>> faces;//定义映射后的人脸
    dlib::matrix<dlib::rgb_pixel> face_chip;
    
    //将人脸坐标其映射到标准人脸上去
    dlib::extract_image_chip(img, dlib::get_face_chip_details(shape, 150, 0.25), face_chip);
    
//     dlib::save_bmp(face_chip, "/data/storage/el2/base/haps/entry/files/face_chip1.bmp");
    faces.push_back(face_chip);
    std::vector<dlib::matrix<float, 0, 1>> face_descriptors = net(faces);//得到128D特征数据
    
    // 创建一个空数组
    napi_value jsArray = nullptr;
    napi_create_array(env, &jsArray);
    // 将创建好的数组进行赋值
    for (int i = 0; i < 128; i++) {
        napi_value element;
        napi_create_double(env, face_descriptors[0](i), &element);
        napi_set_element(env, jsArray, i, element);
    }
    // 返回已创建好的数组
    return jsArray;
}

static napi_value NAPI_Global_detect(napi_env env, napi_callback_info info) {
    
    // 获取并解析传参
    size_t argc = 1;
    napi_value args[1] = {nullptr};
    napi_get_cb_info(env, info, &argc, args , nullptr, nullptr);
    
        // 初始化NativePixelMap对象。
    NativePixelMap *native = OH_PixelMap_InitNativePixelMap(env, args[0]);
    if (native == nullptr) {
        return nullptr;
    }

    // 获取图片信息。
    struct OhosPixelMapInfos pixelMapInfo;
    OH_PixelMap_GetImageInfo(native, &pixelMapInfo);

    // 获取PixelMap对象每行字节数。
    int32_t rowBytes;
    OH_PixelMap_GetBytesNumberPerRow(native, &rowBytes);

    // 获取PixelMap对象是否可编辑的状态。
    int32_t editable = 0;
    OH_PixelMap_GetIsEditable(native, &editable);

    // 获取PixelMap对象是否支持Alpha通道。
    int32_t supportAlpha = 0;
    OH_PixelMap_IsSupportAlpha(native, &supportAlpha);
    
    // 获取PixelMap对象数据的内存地址，并锁定该内存。
    void *pixelAddr = nullptr;
    OH_PixelMap_AccessPixels(native, &pixelAddr);
    
//     array2d<unsigned char> img;
//     rgba_to_dlib_grayscale((unsigned char *)pixelAddr, pixelMapInfo.width, pixelMapInfo.height, img);
    unsigned char *bgr_img = (unsigned char *)malloc(pixelMapInfo.width * pixelMapInfo.height * 3);
    rgba_to_bgr((unsigned char *)pixelAddr, bgr_img, pixelMapInfo.width, pixelMapInfo.height);
    
    // 释放PixelMap对象数据的内存锁。
    OH_PixelMap_UnAccessPixels(native);
    
    int * pResults = NULL; 
    //pBuffer is used in the detection functions.
    //If you call functions in multiple threads, please create one buffer for each thread!
    unsigned char * pBuffer = (unsigned char *)malloc(DETECT_BUFFER_SIZE);
    if(!pBuffer)
    {
//         fprintf(stderr, "Can not alloc buffer.\n");
        return nullptr;
    }
    
    pResults = facedetect_cnn(pBuffer, bgr_img, pixelMapInfo.width, pixelMapInfo.height, pixelMapInfo.width * 3);
    
//     load_image(img, "/data/storage/el2/base/haps/entry/files/3.bmp");
//     dlib::save_bmp(img, "/data/storage/el2/base/haps/entry/files/from_opencv.bmp");
    
//     std::vector<dlib::rectangle> dets = detector(img);
    
    // 创建一个空数组
    napi_value jsArray = nullptr;
    napi_create_array(env, &jsArray);
    
    // 将创建好的数组进行赋值
    int face_number = pResults ? *pResults : 0;//检测到的人脸数，每个区域只有一个人脸
    for (int i = 0; i < face_number; i++) {
//         dlib::rectangle rect = dets[i];
        
        short * p = ((short*)(pResults + 1)) + 16*i;
		int confidence = p[0];
		int x = p[1];
		int y = p[2];
		int w = p[3];
		int h = p[4];
        
        
        napi_value object = nullptr;
        // 创建一个空对象
        napi_create_object(env, &object);
        
        // 设置对象的属性
        napi_value left = nullptr;
        napi_create_string_utf8(env, "left", NAPI_AUTO_LENGTH, &left);
        napi_value left_v = nullptr;
        napi_create_int32(env, x, &left_v);
        
        // 将属性设置到对象上
        napi_set_property(env, object, left, left_v);
        
        // 设置对象的属性
        napi_value top = nullptr;
        napi_create_string_utf8(env, "top", NAPI_AUTO_LENGTH, &top);
        napi_value top_v = nullptr;
        napi_create_int32(env, y, &top_v);
        
        // 将属性设置到对象上
        napi_set_property(env, object, top, top_v);
        
        // 设置对象的属性
        napi_value right = nullptr;
        napi_create_string_utf8(env, "right", NAPI_AUTO_LENGTH, &right);
        napi_value right_v = nullptr;
        napi_create_int32(env, x + w, &right_v);
        
        // 将属性设置到对象上
        napi_set_property(env, object, right, right_v);
        
        // 设置对象的属性
        napi_value bottom = nullptr;
        napi_create_string_utf8(env, "bottom", NAPI_AUTO_LENGTH, &bottom);
        napi_value bottom_v = nullptr;
        napi_create_int32(env, y + h, &bottom_v);
        
        // 将属性设置到对象上
        napi_set_property(env, object, bottom, bottom_v);
        
        napi_set_element(env, jsArray, i, object);
    }
    
    free(bgr_img);
    //release the buffer
    free(pBuffer);
    
    // 返回已创建好的数组
    return jsArray;
}

static napi_value NAPI_Global_initFace(napi_env env, napi_callback_info info) {
//     detector = get_frontal_face_detector();
//     deserialize("/data/storage/el2/base/haps/entry/files/shape_predictor_5_face_landmarks.dat") >> sp;
//     deserialize("/data/storage/el2/base/haps/entry/files/dlib_face_recognition_resnet_model_v1.dat") >> net;
    
    size_t argc = 2;
    napi_value args[2] = {nullptr};

    napi_get_cb_info(env, info, &argc, args, nullptr, nullptr);
    // 获取字符串的长度
    size_t length1 = 0;
    napi_get_value_string_utf8(env, args[0], nullptr, 0, &length1);

    char* buf1 = new char[length1 + 1];
    std::memset(buf1, 0, length1 + 1);
    napi_get_value_string_utf8(env, args[0], buf1, length1 + 1, &length1);
    
    size_t length2 = 0;
    napi_get_value_string_utf8(env, args[1], nullptr, 0, &length2);

    char* buf2 = new char[length2 + 1];
    std::memset(buf2, 0, length2 + 1);
    napi_get_value_string_utf8(env, args[1], buf2, length2 + 1, &length2);
    
    OH_LOG_INFO(LOG_APP, "sp = %{public}s, net = %{public}s", buf1, buf2);
    
    deserialize(buf1) >> sp;
    deserialize(buf2) >> net;
    
    delete[] buf1;
    delete[] buf2;
    
    napi_value ret;
    napi_create_double(env, 0, &ret);

    return ret;
}


static napi_value NAPI_Global_compare(napi_env env, napi_callback_info info) {
    // 获取ArkTS侧传入的参数
    size_t argc = 2;
    napi_value args[2] = {nullptr};

    uint32_t length;
    napi_get_cb_info(env, info, &argc, args, nullptr, nullptr);
    // 检查参数是否为数组
    bool is_array;
    napi_is_array(env, args[0], &is_array);
    if (!is_array) {
        return nullptr;
    }
    napi_get_array_length(env, args[0], &length);
    if (length != 128) {
        return nullptr;
    }
    
    napi_is_array(env, args[1], &is_array);
    if (!is_array) {
        return nullptr;
    }
    napi_get_array_length(env, args[1], &length);
    if (length != 128) {
        return nullptr;
    }
    
    dlib::matrix<float, 0, 1> features1;
    dlib::matrix<float, 0, 1> features2;
    
    features1.set_size(128);
    features2.set_size(128);
    
    for (int i = 0; i < 128; i++) {
        // 获取请求索引位置的元素值并存储在result中
        napi_value result1;
        napi_get_element(env, args[0], i, &result1);
        
        napi_value result2;
        napi_get_element(env, args[1], i, &result2);
        
        double d1;
        napi_get_value_double(env, result1, &d1);
        double d2;
        napi_get_value_double(env, result2, &d2);
        
        features1(i) = d1;
        features2(i) = d2;
    }
    
    float dist = dlib::length(features1 - features2);
    
    double score = distanceToScoreNonlinear(dist);
    
    napi_value ret;
    napi_create_double(env, score, &ret);

    return ret;
}

EXTERN_C_START
static napi_value Init(napi_env env, napi_value exports) {
    napi_property_descriptor desc[] = {
        {"detect", nullptr, NAPI_Global_detect, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"initFace", nullptr, NAPI_Global_initFace, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"feature", nullptr, NAPI_Global_feature, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"compare", nullptr, NAPI_Global_compare, nullptr, nullptr, nullptr, napi_default, nullptr }
    };
    napi_define_properties(env, exports, sizeof(desc) / sizeof(desc[0]), desc);
    return exports;
}
EXTERN_C_END

static napi_module demoModule = {
    .nm_version = 1,
    .nm_flags = 0,
    .nm_filename = nullptr,
    .nm_register_func = Init,
    .nm_modname = "face",
    .nm_priv = ((void*)0),
    .reserved = { 0 },
};

extern "C" __attribute__((constructor)) void RegisterEntryModule(void)
{
    napi_module_register(&demoModule);
}
